#  -*- coding: utf-8 -*-

from pymongo import UpdateOne,ASCENDING, DESCENDING
from factor.base_factor import BaseFactor
from data.finance_report_crawler import FinanceReportCrawler
from data.data_module import DataModule
from util.stock_util import get_all_codes,get_all_indexes_date,multi_computer
from util.database import DB_CONN
import time


"""
实现当日涨跌幅的因子计算和保存
"""


class ChangeRateFactor(BaseFactor):
    def __init__(self):
        BaseFactor.__init__(self, name='change_rate')

    def computer_single(self, code,is_index,begin_date, end_date):

        dm = DataModule()
        start_time = time.time()
        print('计算涨跌幅和是否涨跌停, %s' % code)
        df_daily = dm.get_k_data(code, index=is_index, autype=None, begin_date=begin_date, end_date=end_date)
        if df_daily.index.size > 0:
            df_daily.set_index(['date'], 1, inplace=True)
            update_requests = []
            for date in df_daily.index:

                open = df_daily.loc[date]['open']
                close = df_daily.loc[date]['close']
                high = df_daily.loc[date]['high']
                low = df_daily.loc[date]['low']
                price_limit = 0
                try:
                    pre_close = df_daily.loc[date]['pre_close']
                    if pre_close > 0 and open > 0 and close > 0:
                        open_change_rate = round(100 * (open - pre_close) / pre_close, 2)
                        close_change_rate = round(100 * (close - pre_close) / pre_close, 2)

                        if close >= round(pre_close * 1.1, 2):
                            if open == high == low == close:
                                price_limit = 2
                            else:
                                price_limit = 1
                        if close <= round(pre_close * 0.9, 2):
                            if open == high == low == close:
                                price_limit = -2
                            else:
                                price_limit = -1

                        update_requests.append(
                            UpdateOne(
                                {'code': code, 'date': date,'index': is_index},
                                {'$set': {'code': code,
                                          'date': date,
                                          'index': is_index,
                                          'price_limit': price_limit,
                                          'open_change_rate': open_change_rate,
                                          'close_change_rate': close_change_rate}},
                                upsert=True))
                except:
                    print('填充涨跌幅时发生错误，股票代码：%s，是否指数：%s, 日期：%s' % (code,is_index, date), flush=True)

            if len(update_requests) > 0:
                update_result = self.collection.bulk_write(update_requests, ordered=False)
                end_time = time.time()
                print('填充涨跌幅，股票代码：%s，是否指数：%s, 插入：%4d条，更新：%4d条,耗时：%.3f 秒' %
                      (code, is_index, update_result.upserted_count, update_result.modified_count, (end_time - start_time)),
                      flush=True)

    def compute(self, begin_date, end_date):
        """
        计算指定时间段内所有股票的该因子的值(当日涨跌幅和是否涨跌停)，并保存到数据库中
        price_limit 涨跌停：1 换手涨停，2 一字板涨停，-1 跌停 -2 一字板跌停 0 不是涨跌停
        :param begin_date:  开始时间
        :param end_date: 结束时间
        """
        self.collection.create_index([('code', 1), ('date', 1),('index',1)])

        #获取所有股票
        codes = get_all_codes()
        args = (False, None, begin_date,end_date)
        multi_computer(computer_codes,codes,args)

        #获取所有指数
        codes = get_all_indexes_date()
        args = (True,None, begin_date,end_date)
        multi_computer(computer_codes,codes,args)


def computer_codes(codes, is_index, autype, begin_date, end_date):
    for code in codes:
        ChangeRateFactor().computer_single(code, is_index, begin_date, end_date)


if __name__ == '__main__':
    # 执行因子的提取任务
    ChangeRateFactor().compute('1990-12-19', '2020-11-13')
